from configs import get_config
from solver import Solver
from data_loader import load_h5, VideoData
from torch.utils.data import DataLoader
import numpy as np
import argparse
import torch
import h5py
import platform

# TODO 两个parser好像冲突了
# parser = argparse.ArgumentParser("Pytorch code for unsupervised video summarization")
# Dataset options
# parser.add_argument('--fold-idx', type=int, default=0, help="fold index (default: 0)")
# parser.add_argument('-m', '--metric', type=str, required=True, choices=['tvsum', 'summe'],
#                     help="evaluation metric ['tvsum', 'summe']")
# args = parser.parse_args()


def prepare_fold(fold_idx, train_dict, lbl_t=1):
    """prepare dataset

    assert 剩余的比例是一折的比例
    :fold_idx: 1, 2, ...,
    :lbl_t: TODO
    :returns: TODO

    """
    dataset = {
        'summe': load_h5('../data/', 'SumMe', lbl_t),
        'youtube': load_h5('../data/', 'Youtube', lbl_t),
        'tvsum': load_h5('../data/', 'TVSum', lbl_t),
        'ovp': load_h5('../data/', 'OVP', lbl_t)
    }
    vdo_num = {'summe': 25, 'youtube': 39, 'tvsum': 50, 'ovp': 50}
    train_fea, train_label, train_id = [], [], []
    test_fea, test_label, test_id = [], [], []
    if True:
        for k, rate in train_dict.items():
            if rate == -1:
                continue
            # TODO 0.2的存储问题
            test_num = vdo_num[k] - int(rate*vdo_num[k])
            rand_idx = fold[k]
            test_id = rand_idx[test_num*(fold_idx-1):fold_idx*test_num]
            train_id = rand_idx[0:test_num*(fold_idx-1)] + rand_idx[fold_idx*test_num:]
            for i in test_id:
                test_fea.append(dataset[k][0][i])
                test_label.append(dataset[k][1][i])
            for i in train_id:
                train_fea.append(dataset[k][0][i])
                train_label.append(dataset[k][1][i])
    else:
        for k, rate in train_dict.items():
            if rate == -1:
                continue
            train_num = int(rate*vdo_num[k])
            rand_idx = np.random.permutation(vdo_num[k])
            for i, vdo_id in enumerate(rand_idx):
                if i < train_num:
                    train_fea.append(dataset[k][0][vdo_id])
                    train_label.append(dataset[k][1][vdo_id])
                    train_id.append(vdo_id)
                else:
                    test_fea.append(dataset[k][0][vdo_id])
                    test_label.append(dataset[k][1][vdo_id])
                    test_id.append(vdo_id)
    return train_fea, train_label, train_id, test_fea, test_label, test_id

def prepare(dataset, fold_id):
    """TODO: Docstring for prepare.

    :fold_id: TODO
    :returns: TODO

    """
    fold = {'summe': [ 5, 24, 22, 17, 18,
                      10,  7, 14,  8, 20,
                      3,  6, 15, 23, 12,
                      9,  0, 11,  2,  1,
                      21, 13, 16, 19,  4],
            'tvsum': [12, 24, 35,  2, 18, 48, 36,  4, 37, 38,
                      17, 14, 45,  8, 42, 31,  5, 33, 28, 10,
                      47,  1, 20, 46, 13,  3, 30, 11,  6, 15,
                      7, 40, 26,  0, 44, 32, 16, 21, 23, 25,
                      34, 22, 49, 41, 43, 29, 27, 19, 39,  9],
            'youtube': [26, 35, 44, 22, 19, 16, 28, 20, -1, -1,
                        13, 30, 39, 18, 38, 32, 48, 40, -1, -1,
                        43, 31, 10, 24, 36, 37, 47, 12, -1, -1,
                        41, 15, 34, 45, 25, 11, 46, 17, -1, -1,
                        14, 27, 42, 33, 49, 29, 23, -1, -1, -1 ],
            'ovp': [4, 21, 26, 45, 9, 1, 37, 35, 33, 36,
                    8, 24, 2, 25, 48, 12, 31, 44, 0, 40,
                    3, 47, 5, 39, 43, 46, 15, 42, 22, 7,
                    10, 34, 41, 17, 6, 29, 49, 13, 14, 32,
                    20, 11, 18, 38, 23, 30, 16, 27, 28, 19]}
    train_dict = {
        'dev': {'summe': -1, 'youtube': -1, 'ovp': -1., 'tvsum': 0.6},
        'tvsum_canonical': {'summe': -1, 'youtube': -1, 'ovp': -1., 'tvsum': 0.8},
        'summe_canonical': {'summe': 0.8, 'youtube': -1, 'ovp': -1., 'tvsum': -1},
        'ovp_canonical': {'summe': -1, 'youtube': -1, 'ovp': 0.8, 'tvsum': -1},
        'ytb_canonical': {'summe': -1, 'youtube': 0.8, 'ovp': -1, 'tvsum': -1},
        'tvsum_aug': {'summe': 1, 'youtube': 1, 'ovp': 1., 'tvsum': 0.8},
        'summe_aug': {'summe': 0.8, 'youtube': 1, 'ovp': 1., 'tvsum': 1},
        'ytb_aug': {'summe': 1, 'youtube': 0.8, 'ovp': 1., 'tvsum': 1},
        'ovp_aug': {'summe': 1, 'youtube': 1, 'ovp': 0.8, 'tvsum': 1},
        'tvsum_transfer': {'summe': 1, 'youtube': 1, 'ovp': 1., 'tvsum': 0},
        'summe_transfer': {'summe': 0, 'youtube': 1, 'ovp': 1., 'tvsum': 1},
        'ovp_transfer': {'summe': 1, 'youtube': 1, 'ovp': 0, 'tvsum': 1},
        'ytb_transfer': {'summe': 1, 'youtube': 0, 'ovp': 1, 'tvsum': 1},
    }[config.dataset]

    root = '/home/zgp/pytorch-vsumm-reinforce/datasets'
    dataset = {
        'summe': h5py.File(f'{root}/eccv16_dataset_summe_google_pool5.h5', 'r'),
        'youtube': h5py.File(f'{root}/eccv16_dataset_youtube_google_pool5.h5', 'r'),
        'tvsum': h5py.File(f'{root}/eccv16_dataset_tvsum_google_pool5.h5', 'r'),
        'ovp': h5py.File(f'{root}/eccv16_dataset_ovp_google_pool5.h5', 'r'),
    }
    vdo_num = {'summe': 25, 'youtube': 39+11, 'tvsum': 50, 'ovp': 50}
    train_fea, train_id, train_label = [], [], []
    test_fea, test_id, test_label = [], [], []
    for k, rate in train_dict.items():
        if rate == -1:
            continue
        # TODO 0.2的存储问题
        # Note transfer时只需要跑第一折, 因为一折的比例是100%
        # assert (1-rate)是一折的比例
        test_num = vdo_num[k] - int(rate*vdo_num[k])
        rand_idx = fold[k]
        test_idx = rand_idx[test_num*(fold_id-1):fold_id*test_num]
        train_idx = rand_idx[0:test_num*(fold_id-1)] + rand_idx[fold_id*test_num:]
        for i in test_idx:
            if i == -1: continue
            vdo_k = f'video_{i+1}'
            test_fea.append(dataset[k][vdo_k]['features'][...])
            test_label.append(dataset[k][vdo_k]['gtscore'][...])
            test_id.append(i)
        for i in train_idx:
            if i == -1: continue
            vdo_k = f'video_{i+1}'
            train_fea.append(dataset[k][vdo_k]['features'][...])
            train_label.append(dataset[k][vdo_k]['gtscore'][...])
            train_id.append(i)

    # FIXME: train_id包含相同的数字, train_id并未被使用
    return train_fea, train_label, train_id, test_fea, test_label, test_id

'''
CUDA_VISIBLE_DEVICES=1 nohup python train.py \
    --dataset tvsum_canonical \
    --shift sup \
    --fold 1 \
    --suffix hs1024_2018_tssssss \
    > tvsum_canonical_sup_hs1024__1_2018.log 2>&1&
'''
if __name__ == '__main__':
    # TODO get_config并没有区分设置train和test, 这不合理
    config = get_config()
    print(config)
    # torch.cuda.current_device()  note: 通过命令行变量传入可见设备, 但是为什么环境变量行不通
    print('kernel version:', platform.version())    # $ uname -v
    print(' torch version:', torch.__version__)
    print('   device name:', torch.cuda.get_device_name(0))
    print('python version:', platform.python_version())
    # print(torch.cuda.device_count())

    # test_config = get_config(mode='test')

    # TODO 通过 train, test split 信息即可从文件获取其他信息
    train_fea, train_label, train_id, test_fea, test_label, test_id = prepare(config.dataset, config.fold)
    # train_fea, train_label, train_id, test_fea, test_label, test_id = prepare_fold(fold_idx, train_dict)

    # config.video_root_dir, test_config.video_root_dir
    train_loader = DataLoader(
        VideoData(data=train_fea, label=train_label, vdo_id=train_id),
        batch_size=1)
    test_loader = DataLoader(
        VideoData(with_name=True, data=test_fea, label=test_label, vdo_id=test_id),
        batch_size=1)

    solver = Solver(config, train_loader, test_loader)

    solver.build()
    solver.train()
